{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2022 @ Shenzhen Bay Laboratory & Peking University & Huawei Technologies Co., Ltd\n",
    "\n",
    "This code is a part of Cybertron package.\n",
    "\n",
    "The Cybertron is open-source software based on the AI-framework:\n",
    "MindSpore (https://www.mindspore.cn/)\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "\n",
    "You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "Cybertron tutorial 04: Read parameters and hyperparameters from checkpoint file\n",
    "                       and to use test dataset with scale and shift."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(14474:140502773348160,MainProcess):2022-08-10-17:32:19.847.768 [mindspore/run_check/_check_version.py:137] Can not found cuda libs, please confirm that the correct cuda version has been installed, you can refer to the installation guidelines: https://www.mindspore.cn/install\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "from mindspore import nn\n",
    "from mindspore import context\n",
    "from mindspore import dataset as ds\n",
    "from mindspore.train import Model\n",
    "from mindspore.train import load_checkpoint\n",
    "\n",
    "from mindsponge.data import load_hyperparam\n",
    "\n",
    "from cybertron import Cybertron\n",
    "from cybertron.train import MAE\n",
    "from cybertron.train import WithLabelEvalCell\n",
    "\n",
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_file = 'Tutorial_C03/Tutorial_C03_MolCT-best.ckpt'\n",
    "hyper_param = load_hyperparam(ckpt_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model.atom_embedding.embedding_table': Parameter (name=model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.linear.weight': Parameter (name=model.dis_filter.linear.weight, shape=(128, 64), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.linear.bias': Parameter (name=model.dis_filter.linear.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.residual.nonlinear.mlp.0.weight': Parameter (name=model.dis_filter.residual.nonlinear.mlp.0.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.residual.nonlinear.mlp.0.bias': Parameter (name=model.dis_filter.residual.nonlinear.mlp.0.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.residual.nonlinear.mlp.1.weight': Parameter (name=model.dis_filter.residual.nonlinear.mlp.1.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.dis_filter.residual.nonlinear.mlp.1.bias': Parameter (name=model.dis_filter.residual.nonlinear.mlp.1.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.positional_embedding.norm.gamma': Parameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.positional_embedding.norm.beta': Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.positional_embedding.x2q.weight': Parameter (name=model.interactions.0.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.positional_embedding.x2k.weight': Parameter (name=model.interactions.0.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.positional_embedding.x2v.weight': Parameter (name=model.interactions.0.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.0.multi_head_attention.output.weight': Parameter (name=model.interactions.0.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.positional_embedding.norm.gamma': Parameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.positional_embedding.norm.beta': Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.positional_embedding.x2q.weight': Parameter (name=model.interactions.1.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.positional_embedding.x2k.weight': Parameter (name=model.interactions.1.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.positional_embedding.x2v.weight': Parameter (name=model.interactions.1.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.1.multi_head_attention.output.weight': Parameter (name=model.interactions.1.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.positional_embedding.norm.gamma': Parameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.positional_embedding.norm.beta': Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.positional_embedding.x2q.weight': Parameter (name=model.interactions.2.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.positional_embedding.x2k.weight': Parameter (name=model.interactions.2.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.positional_embedding.x2v.weight': Parameter (name=model.interactions.2.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'model.interactions.2.multi_head_attention.output.weight': Parameter (name=model.interactions.2.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'readout.decoder.output.mlp.0.weight': Parameter (name=readout.decoder.output.mlp.0.weight, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'readout.decoder.output.mlp.0.bias': Parameter (name=readout.decoder.output.mlp.0.bias, shape=(64,), dtype=Float32, requires_grad=True),\n",
       " 'readout.decoder.output.mlp.1.weight': Parameter (name=readout.decoder.output.mlp.1.weight, shape=(1, 64), dtype=Float32, requires_grad=True),\n",
       " 'readout.decoder.output.mlp.1.bias': Parameter (name=readout.decoder.output.mlp.1.bias, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'global_step': Parameter (name=global_step, shape=(1,), dtype=Int32, requires_grad=True),\n",
       " 'beta1_power': Parameter (name=beta1_power, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'beta2_power': Parameter (name=beta2_power, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.atom_embedding.embedding_table': Parameter (name=moment1.model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.linear.weight': Parameter (name=moment1.model.dis_filter.linear.weight, shape=(128, 64), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.linear.bias': Parameter (name=moment1.model.dis_filter.linear.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.residual.nonlinear.mlp.0.weight': Parameter (name=moment1.model.dis_filter.residual.nonlinear.mlp.0.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.residual.nonlinear.mlp.0.bias': Parameter (name=moment1.model.dis_filter.residual.nonlinear.mlp.0.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.residual.nonlinear.mlp.1.weight': Parameter (name=moment1.model.dis_filter.residual.nonlinear.mlp.1.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.dis_filter.residual.nonlinear.mlp.1.bias': Parameter (name=moment1.model.dis_filter.residual.nonlinear.mlp.1.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.positional_embedding.norm.gamma': Parameter (name=moment1.model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.positional_embedding.norm.beta': Parameter (name=moment1.model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.positional_embedding.x2q.weight': Parameter (name=moment1.model.interactions.0.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.positional_embedding.x2k.weight': Parameter (name=moment1.model.interactions.0.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.positional_embedding.x2v.weight': Parameter (name=moment1.model.interactions.0.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.0.multi_head_attention.output.weight': Parameter (name=moment1.model.interactions.0.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.positional_embedding.norm.gamma': Parameter (name=moment1.model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.positional_embedding.norm.beta': Parameter (name=moment1.model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.positional_embedding.x2q.weight': Parameter (name=moment1.model.interactions.1.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.positional_embedding.x2k.weight': Parameter (name=moment1.model.interactions.1.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.positional_embedding.x2v.weight': Parameter (name=moment1.model.interactions.1.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.1.multi_head_attention.output.weight': Parameter (name=moment1.model.interactions.1.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.positional_embedding.norm.gamma': Parameter (name=moment1.model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.positional_embedding.norm.beta': Parameter (name=moment1.model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.positional_embedding.x2q.weight': Parameter (name=moment1.model.interactions.2.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.positional_embedding.x2k.weight': Parameter (name=moment1.model.interactions.2.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.positional_embedding.x2v.weight': Parameter (name=moment1.model.interactions.2.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.model.interactions.2.multi_head_attention.output.weight': Parameter (name=moment1.model.interactions.2.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.readout.decoder.output.mlp.0.weight': Parameter (name=moment1.readout.decoder.output.mlp.0.weight, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment1.readout.decoder.output.mlp.0.bias': Parameter (name=moment1.readout.decoder.output.mlp.0.bias, shape=(64,), dtype=Float32, requires_grad=True),\n",
       " 'moment1.readout.decoder.output.mlp.1.weight': Parameter (name=moment1.readout.decoder.output.mlp.1.weight, shape=(1, 64), dtype=Float32, requires_grad=True),\n",
       " 'moment1.readout.decoder.output.mlp.1.bias': Parameter (name=moment1.readout.decoder.output.mlp.1.bias, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.atom_embedding.embedding_table': Parameter (name=moment2.model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.linear.weight': Parameter (name=moment2.model.dis_filter.linear.weight, shape=(128, 64), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.linear.bias': Parameter (name=moment2.model.dis_filter.linear.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.residual.nonlinear.mlp.0.weight': Parameter (name=moment2.model.dis_filter.residual.nonlinear.mlp.0.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.residual.nonlinear.mlp.0.bias': Parameter (name=moment2.model.dis_filter.residual.nonlinear.mlp.0.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.residual.nonlinear.mlp.1.weight': Parameter (name=moment2.model.dis_filter.residual.nonlinear.mlp.1.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.dis_filter.residual.nonlinear.mlp.1.bias': Parameter (name=moment2.model.dis_filter.residual.nonlinear.mlp.1.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.positional_embedding.norm.gamma': Parameter (name=moment2.model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.positional_embedding.norm.beta': Parameter (name=moment2.model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.positional_embedding.x2q.weight': Parameter (name=moment2.model.interactions.0.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.positional_embedding.x2k.weight': Parameter (name=moment2.model.interactions.0.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.positional_embedding.x2v.weight': Parameter (name=moment2.model.interactions.0.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.0.multi_head_attention.output.weight': Parameter (name=moment2.model.interactions.0.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.positional_embedding.norm.gamma': Parameter (name=moment2.model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.positional_embedding.norm.beta': Parameter (name=moment2.model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.positional_embedding.x2q.weight': Parameter (name=moment2.model.interactions.1.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.positional_embedding.x2k.weight': Parameter (name=moment2.model.interactions.1.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.positional_embedding.x2v.weight': Parameter (name=moment2.model.interactions.1.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.1.multi_head_attention.output.weight': Parameter (name=moment2.model.interactions.1.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.positional_embedding.norm.gamma': Parameter (name=moment2.model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.positional_embedding.norm.beta': Parameter (name=moment2.model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.positional_embedding.x2q.weight': Parameter (name=moment2.model.interactions.2.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.positional_embedding.x2k.weight': Parameter (name=moment2.model.interactions.2.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.positional_embedding.x2v.weight': Parameter (name=moment2.model.interactions.2.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.model.interactions.2.multi_head_attention.output.weight': Parameter (name=moment2.model.interactions.2.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.readout.decoder.output.mlp.0.weight': Parameter (name=moment2.readout.decoder.output.mlp.0.weight, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'moment2.readout.decoder.output.mlp.0.bias': Parameter (name=moment2.readout.decoder.output.mlp.0.bias, shape=(64,), dtype=Float32, requires_grad=True),\n",
       " 'moment2.readout.decoder.output.mlp.1.weight': Parameter (name=moment2.readout.decoder.output.mlp.1.weight, shape=(1, 64), dtype=Float32, requires_grad=True),\n",
       " 'moment2.readout.decoder.output.mlp.1.bias': Parameter (name=moment2.readout.decoder.output.mlp.1.bias, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.atom_embedding.embedding_table': Parameter (name=vhat.model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.linear.weight': Parameter (name=vhat.model.dis_filter.linear.weight, shape=(128, 64), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.linear.bias': Parameter (name=vhat.model.dis_filter.linear.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.residual.nonlinear.mlp.0.weight': Parameter (name=vhat.model.dis_filter.residual.nonlinear.mlp.0.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.residual.nonlinear.mlp.0.bias': Parameter (name=vhat.model.dis_filter.residual.nonlinear.mlp.0.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.residual.nonlinear.mlp.1.weight': Parameter (name=vhat.model.dis_filter.residual.nonlinear.mlp.1.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.dis_filter.residual.nonlinear.mlp.1.bias': Parameter (name=vhat.model.dis_filter.residual.nonlinear.mlp.1.bias, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.positional_embedding.norm.gamma': Parameter (name=vhat.model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.positional_embedding.norm.beta': Parameter (name=vhat.model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.positional_embedding.x2q.weight': Parameter (name=vhat.model.interactions.0.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.positional_embedding.x2k.weight': Parameter (name=vhat.model.interactions.0.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.positional_embedding.x2v.weight': Parameter (name=vhat.model.interactions.0.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.0.multi_head_attention.output.weight': Parameter (name=vhat.model.interactions.0.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.positional_embedding.norm.gamma': Parameter (name=vhat.model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.positional_embedding.norm.beta': Parameter (name=vhat.model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.positional_embedding.x2q.weight': Parameter (name=vhat.model.interactions.1.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.positional_embedding.x2k.weight': Parameter (name=vhat.model.interactions.1.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.positional_embedding.x2v.weight': Parameter (name=vhat.model.interactions.1.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.1.multi_head_attention.output.weight': Parameter (name=vhat.model.interactions.1.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.positional_embedding.norm.gamma': Parameter (name=vhat.model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.positional_embedding.norm.beta': Parameter (name=vhat.model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.positional_embedding.x2q.weight': Parameter (name=vhat.model.interactions.2.positional_embedding.x2q.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.positional_embedding.x2k.weight': Parameter (name=vhat.model.interactions.2.positional_embedding.x2k.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.positional_embedding.x2v.weight': Parameter (name=vhat.model.interactions.2.positional_embedding.x2v.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.model.interactions.2.multi_head_attention.output.weight': Parameter (name=vhat.model.interactions.2.multi_head_attention.output.weight, shape=(128, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.readout.decoder.output.mlp.0.weight': Parameter (name=vhat.readout.decoder.output.mlp.0.weight, shape=(64, 128), dtype=Float32, requires_grad=True),\n",
       " 'vhat.readout.decoder.output.mlp.0.bias': Parameter (name=vhat.readout.decoder.output.mlp.0.bias, shape=(64,), dtype=Float32, requires_grad=True),\n",
       " 'vhat.readout.decoder.output.mlp.1.weight': Parameter (name=vhat.readout.decoder.output.mlp.1.weight, shape=(1, 64), dtype=Float32, requires_grad=True),\n",
       " 'vhat.readout.decoder.output.mlp.1.bias': Parameter (name=vhat.readout.decoder.output.mlp.1.bias, shape=(1,), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.name': Parameter (name=hyperparam.model.name, shape=(5,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.num_atom_types': Parameter (name=hyperparam.model.num_atom_types, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.num_bond_types': Parameter (name=hyperparam.model.num_bond_types, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.dim_feature': Parameter (name=hyperparam.model.dim_feature, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.n_interaction': Parameter (name=hyperparam.model.n_interaction, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.activation.name': Parameter (name=hyperparam.model.activation.name, shape=(5,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.cutoff': Parameter (name=hyperparam.model.cutoff, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.cutoff_fn.name': Parameter (name=hyperparam.model.cutoff_fn.name, shape=(12,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.rbf.name': Parameter (name=hyperparam.model.rbf.name, shape=(16,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.rbf.r_max': Parameter (name=hyperparam.model.rbf.r_max, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.rbf.r_min': Parameter (name=hyperparam.model.rbf.r_min, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.rbf.sigma': Parameter (name=hyperparam.model.rbf.sigma, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.rbf.delta': Parameter (name=hyperparam.model.rbf.delta, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.rbf.num_basis': Parameter (name=hyperparam.model.rbf.num_basis, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.rbf.clip_distance': Parameter (name=hyperparam.model.rbf.clip_distance, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.rbf.rescale': Parameter (name=hyperparam.model.rbf.rescale, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.rbf.length_unit': Parameter (name=hyperparam.model.rbf.length_unit, shape=(2,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.rbf.r_ref': Parameter (name=hyperparam.model.rbf.r_ref, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.r_self': Parameter (name=hyperparam.model.r_self, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.model.coupled_interaction': Parameter (name=hyperparam.model.coupled_interaction, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.use_distance': Parameter (name=hyperparam.model.use_distance, shape=(3,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.use_bond': Parameter (name=hyperparam.model.use_bond, shape=(3,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.public_dis_filter': Parameter (name=hyperparam.model.public_dis_filter, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.public_bond_filter': Parameter (name=hyperparam.model.public_bond_filter, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.use_graph_norm': Parameter (name=hyperparam.model.use_graph_norm, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.length_unit': Parameter (name=hyperparam.model.length_unit, shape=(2,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.model.n_heads': Parameter (name=hyperparam.model.n_heads, shape=(3,), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.max_cycles': Parameter (name=hyperparam.model.max_cycles, shape=(3,), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.model.use_feed_forward': Parameter (name=hyperparam.model.use_feed_forward, shape=(3,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.fixed_cycles': Parameter (name=hyperparam.model.fixed_cycles, shape=(3,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.model.act_threshold': Parameter (name=hyperparam.model.act_threshold, shape=(3,), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.num_readout': Parameter (name=hyperparam.num_readout, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.readout.name': Parameter (name=hyperparam.readout.name, shape=(15,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.readout.dim_represent': Parameter (name=hyperparam.readout.dim_represent, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.readout.dim_output': Parameter (name=hyperparam.readout.dim_output, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.readout.activation.name': Parameter (name=hyperparam.readout.activation.name, shape=(5,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.readout.decoder.name': Parameter (name=hyperparam.readout.decoder.name, shape=(12,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.readout.aggregator.name': Parameter (name=hyperparam.readout.aggregator.name, shape=(15,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.readout.scale': Parameter (name=hyperparam.readout.scale, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.readout.shift': Parameter (name=hyperparam.readout.shift, shape=(), dtype=Float32, requires_grad=True),\n",
       " 'hyperparam.readout.atomwise_scaleshift': Parameter (name=hyperparam.readout.atomwise_scaleshift, shape=(1,), dtype=Bool, requires_grad=True),\n",
       " 'hyperparam.readout.axis': Parameter (name=hyperparam.readout.axis, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.readout.n_decoder_layers': Parameter (name=hyperparam.readout.n_decoder_layers, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.dim_output': Parameter (name=hyperparam.dim_output, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.num_atoms': Parameter (name=hyperparam.num_atoms, shape=(), dtype=Int64, requires_grad=True),\n",
       " 'hyperparam.length_unit': Parameter (name=hyperparam.length_unit, shape=(2,), dtype=Int8, requires_grad=True),\n",
       " 'hyperparam.energy_unit': Parameter (name=hyperparam.energy_unit, shape=(4,), dtype=Int8, requires_grad=True)}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = Cybertron(hyper_param=hyper_param)\n",
    "load_checkpoint(ckpt_file, net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "Cybertron Engine, Ride-on!\n",
      "--------------------------------------------------------------------------------\n",
      "    Length unit: nm\n",
      "    Input unit scale: 1\n",
      "--------------------------------------------------------------------------------\n",
      "    Deep molecular model:  MolCT\n",
      "--------------------------------------------------------------------------------\n",
      "       Length unit: nm\n",
      "       Atom embedding size: 64\n",
      "       Cutoff distance: 1.0 nm\n",
      "       Radical basis function (RBF): LogGaussianBasis\n",
      "          Minimum distance: 0.04 nm\n",
      "          Maximum distance: 1.0 nm\n",
      "          Reference distance: 1.0 nm\n",
      "          Log Gaussian begin: -3.218876\n",
      "          Log Gaussian end: 0.006724119\n",
      "          Interval for log Gaussian: 0.0512\n",
      "          Sigma for log gaussian: 0.3\n",
      "          Number of basis functions: 64\n",
      "          Rescale the range of RBF to (-1,1).\n",
      "       Calculate distance: Yes\n",
      "       Calculate bond: No\n",
      "       Feature dimension: 128\n",
      "--------------------------------------------------------------------------------\n",
      "       Using 3 independent interaction layers:\n",
      "--------------------------------------------------------------------------------\n",
      "       0. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       1. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       2. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "    Readout network: AtomwiseReadout\n",
      "--------------------------------------------------------------------------------\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: [18.248537]\n",
      "       Shift: [-409.42038]\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types:\n",
      "          No.0:   [0.]\n",
      "          No.1:   [-1313.4669]\n",
      "          No.2:   [0.]\n",
      "          No.3:   [0.]\n",
      "          No.4:   [0.]\n",
      "          No.5:   [0.]\n",
      "          No.6:   [-99366.71]\n",
      "          No.7:   [-143309.94]\n",
      "          No.8:   [-197082.06]\n",
      "          No.9:   [-261811.55]\n",
      "       Output unit: kJ mol-1\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    Output dimension: 1\n",
      "    Output unit for Cybertron: kJ mol-1\n",
      "    Output unit scale: 1.0\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "idx = [7]  # U0\n",
    "\n",
    "test_file = sys.path[0] + '/dataset_qm9_origin_testset_1024.npz'\n",
    "test_data = np.load(test_file)\n",
    "\n",
    "scale = test_data['scale'][idx]\n",
    "shift = test_data['shift'][idx]\n",
    "ref = test_data['type_ref'][:, idx]\n",
    "\n",
    "net.set_scaleshift(scale=scale, shift=shift, type_ref=ref, unit='kj/mol')\n",
    "net.print_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_test = ds.NumpySlicesDataset(\n",
    "    {'R': test_data['R'], 'Z': test_data['Z'], 'E': test_data['E'][:, idx]}, shuffle=False)\n",
    "ds_test = ds_test.batch(1024)\n",
    "ds_test = ds_test.repeat(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WithLabelEvalCell with input type: RZE\n"
     ]
    }
   ],
   "source": [
    "eval_network = WithLabelEvalCell('RZE', net, nn.MAELoss())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_mae = 'EvalMAE'\n",
    "atom_mae = 'AtomMAE'\n",
    "model = Model(net, eval_network=eval_network, metrics={eval_mae: MAE(\n",
    "    [1, 2]), atom_mae: MAE([1, 2, 3], averaged_by_atoms=True)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test dataset:\n",
      "EvalMAE: 43.42962646484375, AtomMAE: 2.5831985473632812, \n"
     ]
    }
   ],
   "source": [
    "print('Test dataset:')\n",
    "eval_metrics = model.eval(ds_test, dataset_sink_mode=False)\n",
    "info = ''\n",
    "for k, v in eval_metrics.items():\n",
    "    info += k\n",
    "    info += ': '\n",
    "    info += str(v)\n",
    "    info += ', '\n",
    "print(info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Cybertron<\n",
       "  (model): MolCT<\n",
       "    (activation): Swish<\n",
       "      (sigmoid): Sigmoid<>\n",
       "      >\n",
       "    (atom_embedding): Embedding<vocab_size=64, embedding_size=128, use_one_hot=True, embedding_table=Parameter (name=model.atom_embedding.embedding_table, shape=(64, 128), dtype=Float32, requires_grad=True), dtype=Float32, padding_idx=None>\n",
       "    (cutoff_fn): SmoothCutoff<>\n",
       "    (rbf): LogGaussianBasis<>\n",
       "    (dis_filter): ResFilter<\n",
       "      (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "      (residual): Residual<\n",
       "        (nonlinear): MLP<\n",
       "          (mlp): SequentialCell<\n",
       "            (0): Dense<\n",
       "              input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "              (activation): Swish<\n",
       "                (sigmoid): Sigmoid<>\n",
       "                >\n",
       "              >\n",
       "            (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "            >\n",
       "          >\n",
       "        >\n",
       "      >\n",
       "    (interactions): CellList<\n",
       "      (0): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.0.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.0.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      (1): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.1.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.1.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      (2): NeuralInteractionUnit<\n",
       "        (activation): Swish<\n",
       "          (sigmoid): Sigmoid<>\n",
       "          >\n",
       "        (dis_filter): ResFilter<\n",
       "          (linear): Dense<input_channels=64, output_channels=128, has_bias=True>\n",
       "          (residual): Residual<\n",
       "            (nonlinear): MLP<\n",
       "              (mlp): SequentialCell<\n",
       "                (0): Dense<\n",
       "                  input_channels=128, output_channels=128, has_bias=True, activation=Swish<>\n",
       "                  (activation): Swish<\n",
       "                    (sigmoid): Sigmoid<>\n",
       "                    >\n",
       "                  >\n",
       "                (1): Dense<input_channels=128, output_channels=128, has_bias=True>\n",
       "                >\n",
       "              >\n",
       "            >\n",
       "          >\n",
       "        (positional_embedding): PositionalEmbedding<\n",
       "          (norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (g_norm): LayerNorm<normalized_shape=(128,), begin_norm_axis=-1, begin_params_axis=-1, gammaParameter (name=model.interactions.2.positional_embedding.norm.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=model.interactions.2.positional_embedding.norm.beta, shape=(128,), dtype=Float32, requires_grad=True)>\n",
       "          (x2q): Dense<input_channels=128, output_channels=128>\n",
       "          (x2k): Dense<input_channels=128, output_channels=128>\n",
       "          (x2v): Dense<input_channels=128, output_channels=128>\n",
       "          >\n",
       "        (multi_head_attention): MultiheadAttention<\n",
       "          (output): Dense<input_channels=128, output_channels=128>\n",
       "          (softmax_with_mask): SoftmaxWithMask<>\n",
       "          >\n",
       "        >\n",
       "      >\n",
       "    >\n",
       "  (activation): Swish<\n",
       "    (sigmoid): Sigmoid<>\n",
       "    >\n",
       "  (readout): AtomwiseReadout<\n",
       "    (activation): Swish<\n",
       "      (sigmoid): Sigmoid<>\n",
       "      >\n",
       "    (decoder): HalveDecoder<\n",
       "      (activation): Swish<\n",
       "        (sigmoid): Sigmoid<>\n",
       "        >\n",
       "      (output): MLP<\n",
       "        (mlp): SequentialCell<\n",
       "          (0): Dense<\n",
       "            input_channels=128, output_channels=64, has_bias=True, activation=Swish<>\n",
       "            (activation): Swish<\n",
       "              (sigmoid): Sigmoid<>\n",
       "              >\n",
       "            >\n",
       "          (1): Dense<input_channels=64, output_channels=1, has_bias=True>\n",
       "          >\n",
       "        >\n",
       "      >\n",
       "    (aggregator): TensorSummation<>\n",
       "    >\n",
       "  (fc_neighbours): FullConnectNeighbours<>\n",
       "  (distances): IndexDistances<\n",
       "    (get_vector): GetVector<>\n",
       "    (norm_last_dim): Norm<axis=-1, keep_dims=False>\n",
       "    >\n",
       "  >"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outdir = 'Tutorial_C04'\n",
    "scaled_ckpt = outdir + '_' + net.model_name + '.ckpt'\n",
    "net.save_checkpoint(scaled_ckpt, outdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "Cybertron Engine, Ride-on!\n",
      "--------------------------------------------------------------------------------\n",
      "    Length unit: nm\n",
      "    Input unit scale: 1\n",
      "--------------------------------------------------------------------------------\n",
      "    Deep molecular model:  MolCT\n",
      "--------------------------------------------------------------------------------\n",
      "       Length unit: nm\n",
      "       Atom embedding size: 64\n",
      "       Cutoff distance: 1.0 nm\n",
      "       Radical basis function (RBF): LogGaussianBasis\n",
      "          Minimum distance: 0.04 nm\n",
      "          Maximum distance: 1.0 nm\n",
      "          Reference distance: 1.0 nm\n",
      "          Log Gaussian begin: -3.218876\n",
      "          Log Gaussian end: 0.006724119\n",
      "          Interval for log Gaussian: 0.0512\n",
      "          Sigma for log gaussian: 0.3\n",
      "          Number of basis functions: 64\n",
      "          Rescale the range of RBF to (-1,1).\n",
      "       Calculate distance: Yes\n",
      "       Calculate bond: No\n",
      "       Feature dimension: 128\n",
      "--------------------------------------------------------------------------------\n",
      "       Using 3 independent interaction layers:\n",
      "--------------------------------------------------------------------------------\n",
      "       0. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       1. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       2. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "    Readout network: AtomwiseReadout\n",
      "--------------------------------------------------------------------------------\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: [18.248537]\n",
      "       Shift: [-409.42038]\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types:\n",
      "          No.0:   [0.]\n",
      "          No.1:   [-1313.4669]\n",
      "          No.2:   [0.]\n",
      "          No.3:   [0.]\n",
      "          No.4:   [0.]\n",
      "          No.5:   [0.]\n",
      "          No.6:   [-99366.71]\n",
      "          No.7:   [-143309.94]\n",
      "          No.8:   [-197082.06]\n",
      "          No.9:   [-261811.55]\n",
      "       Output unit: kJ mol-1\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    Output dimension: 1\n",
      "    Output unit for Cybertron: kJ mol-1\n",
      "    Output unit scale: 1.0\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "net2 = Cybertron(hyper_param=load_hyperparam(outdir+'/'+scaled_ckpt))\n",
    "net2.print_info()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 ('mindsponge')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2496ecc683137a232cae2452fbbdd53dab340598b6e499c8995be760f3a431b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
